#define ZCORE_SOURCE
#include "CmdParser.hpp"
#include "StringTools.hpp"
#include "Tools.hpp"
#include "FileTools.hpp"
#include "Log.hpp"
#include <yaml-cpp/yaml.h>

namespace zzz{
ZFLAGS_SWITCH2(help, "Show Help", 'h');
ZFLAGS_SWITCH(helplong, "Show All Help");

ZFLAGS_STRING(opt_file, "", "Option filename (If it is empty, will use program_name.zopt)");
ZFLAGS_STRING(opt_dir, GetEnv("ZOPT_DIR"), "Will find option file in initial folder in current folder or this folder");

const int INDEX_BEGIN=10000;

ZGLOBAL_DEFINE(string, zCurrentProgramName, "");

void CmdParser::AddLongOption(const string &longopt, ArgRequirement re, int shortopt, const string &desc, const string &file) {
  option o;
  o.name=new char[longopt.size()+1];
  strcpy((char*)o.name,longopt.c_str());
  const int argre[3]={no_argument,optional_argument,required_argument};
  o.has_arg=argre[re];
  o.flag=0;
  o.val=shortopt;

  long_options.push_back(o);

  if (shortopt<255 && (isalpha(shortopt) || isdigit(shortopt))) {
    short_options+=shortopt;
    optstr.push_back(string("-") + char(shortopt) + "|");
    if (re==CMDPARSER_ONEARG) short_options+=':';
    else if (re==CMDPARSER_OPTARG) short_options+="::";
  }
  else
    optstr.push_back("");

  optstr.back()+=string("--")+longopt;
  optdesc.push_back(string(desc));

  if (!option_file.empty() && option_file.back()==file) {
    option_file_n.back()++;
  } else {
    option_file.push_back(file);
    option_file_n.push_back(1);
  }

}

void CmdParser::AddShortOption(char shortopt, ArgRequirement re, const string &desc, const string &file) {
  short_options+=shortopt;
  optstr.push_back("-");
  optstr.back()+=shortopt;

  if (re==CMDPARSER_ONEARG) {
    short_options+=':';
  } else if (re==CMDPARSER_OPTARG) {
    short_options+="::";
  }
  optdesc.push_back(string(desc));

  if (option_file.back()==file) {
    option_file_n.back()++;
  }else {
    option_file.push_back(file);
    option_file_n.push_back(1);
  }

}

bool CmdParser::Parse(const string &usage, int argc, char *argv[], const char *currentfile)
{
  *ZGLOBAL_GET(string*, "zCurrentProgramName")=argv[0];

  LoadOptionFile();

  sequential_.clear();
  optind=0;

  if (long_options.back().name!=NULL) {
    option o;
    o.name=NULL;
    o.has_arg=no_argument;
    o.flag=0;
    o.val=0;
    long_options.push_back(o);
  }

  while(true) {
    int option_index;
    int c=getopt_long(argc,argv,short_options.c_str(), &(long_options[0]), &option_index);
    if (c==-1) break;  //end
    if (c==0) continue;  //flag
    if (c=='?') {
      zout<<"Please read " << argv[0] << " --help OR --helplong\n";
      exit(-1);  //error
    }
    if (ParseShortOption(c, optarg)) continue;
    ZLOG(ZFATAL)<<"Unknown opt: " << c << "\n[THIS SHOULD NOT HAPPEN!]";
  }

  if (optind < argc) {
    while (optind < argc)
      sequential_.push_back(string(argv[optind++]));
  }

  if (ZFLAG_help==true) {
    Help(usage, currentfile);
    exit(0);
  }

  if (ZFLAG_helplong==true) {
    HelpLong(usage);
    exit(0);
  }

  for (zuint i=0; i<postparse_.size(); ++i) {
    (postparse_[i])();
  }
  return true;
}

void CmdParser::Help(const string &usage, const char *currentfile) {
  // usage
  cout << "I was built on: " << ZGLOBAL_GET(string, "__MAIN_TIMESTAMP__") << endl;
  cout << "USAGE: " << usage << endl << endl;
  zuint start_opt = 0;
  zuint end_opt = 0;
  for (zuint i = 0; i < option_file.size(); ++i) {
    if (option_file[i] == currentfile) {
      start_opt = accumulate(option_file_n.begin(), option_file_n.begin() + i, 0);
      end_opt = start_opt + option_file_n[i];
      break;
    }
  }
  cout << "OPTIONS in " << currentfile << ":\n";
  // option description
  for (zuint i = start_opt; i < end_opt; ++i) {
    cout.width(18);
    cout << left << optstr[i];
    cout.width(60);
    cout << left << optdesc[i] << endl;
  }
}

void CmdParser::HelpLong(const string &usage)
{
  //usage
  cout << "USAGE: " << usage << endl << endl;
  zuint i=0,iend=0;
  for (zuint j=0; j<option_file.size(); j++) {
    cout<<"OPTIONS in "<<option_file[j]<<":\n";
    //option description
    iend+=option_file_n[j];
    for (; i<iend; ++i) {
      cout.width(18);
      cout << left << optstr[i];
      cout.width(60);
      cout << left << optdesc[i] << endl;
    }
    cout << endl;
  }
}

CmdParser::~CmdParser()
{
  if (long_options.back().name==NULL)
  {
    for (zuint i=0; i<long_options.size()-1; ++i)
      delete[] long_options[i].name;
  }
  else
  {
    for (zuint i=0; i<long_options.size(); ++i)
      delete[] long_options[i].name;
  }
}

void CmdParser::AddAutoOption(const string &longopt, const string &desc, const string &file, const char shortopt, bool is_switch)
{
  string strvar = string("ZFLAG_") + longopt;
  string strdesc(desc);
  if (ZGLOBAL_ISTYPE(double*, strvar)) {
    strdesc += "(" + ToString(ZFLAG_GET(double, strvar)) + ")";
  } else if (ZGLOBAL_ISTYPE(int*, strvar)) {
    strdesc += "(" + ToString(ZFLAG_GET(int, strvar)) + ")";
  } else if (ZGLOBAL_ISTYPE(string*, strvar)) {
    strdesc += "(\"" + ZFLAG_GET(string, strvar) + "\")";
  } else if (ZGLOBAL_ISTYPE(vector<string>*, strvar)) {
    strdesc << "(" << GetzVar(ZFLAG_GET(vector<string>, strvar)) << ")";
  } else if (ZGLOBAL_ISTYPE(bool*, strvar) && !is_switch) {
    strdesc += "(" + ToString(ZFLAG_GET(bool, strvar)) + ")";
  }

  int index;
  if (shortopt==0)  // No short option, assign a number.
    index = INDEX_BEGIN + short_long_.size() + neg_short_long_.size();
  else
    index = shortopt;
  if (ZGLOBAL_ISTYPE(bool*, strvar)) {
    if (is_switch)
      AddLongOption(longopt, CMDPARSER_NOARG, index, strdesc.c_str(), file);
    else
      AddLongOption(longopt, CMDPARSER_OPTARG, index, strdesc.c_str(), file);
  } else {
    AddLongOption(longopt, CMDPARSER_ONEARG, index, strdesc.c_str(), file);
  }

  short_long_[index] = strvar;

  if (!is_switch && ZGLOBAL_ISTYPE(bool*, strvar)) {
    int neg_index = INDEX_BEGIN + short_long_.size() + neg_short_long_.size();
    string negopt = string("no_") + longopt;
    AddLongOption(negopt.c_str(), CMDPARSER_NOARG, neg_index, "", file);
    neg_short_long_[neg_index] = strvar;
  }
}

void CmdParser::AddPostParseFunc(PostParse post)
{
  postparse_.push_back(post);
}

bool CmdParser::ParseShortOption(const int c, const char *arg)
{
  map<int, string>::const_iterator mi = short_long_.find(c);
  if (mi == short_long_.end()) {
    mi = neg_short_long_.find(c);
    if (mi == neg_short_long_.end())
      return false;
    return ParseLongOption(mi->second, "false");
  }
  return ParseLongOption(mi->second, arg);
}

bool CmdParser::ParseLongOption(const string &name, const char *arg) {
  if (!ZGLOBAL_ISEXIST(name))
    return false;

  if (ZGLOBAL_ISTYPE(bool*, name)) {
    if (arg == NULL)
      *(ZGLOBAL_GET(bool*,name)) = true;
    else
      *(ZGLOBAL_GET(bool*,name)) = FromString<bool>(arg);
    return true;
  }

  ZCHECK_NOT_NULL(arg) << ZVAR(name);

#define TRY_TYPE(type) \
  if (ZGLOBAL_ISTYPE(type*, name)) {*(ZGLOBAL_GET(type*,name)) = FromString<type>(arg); return true; }

  TRY_TYPE(double);
  TRY_TYPE(int);
  TRY_TYPE(std::string);

#undef TRY_TYPE


  if (ZGLOBAL_ISTYPE(std::vector<std::string>*, name)) {
    ZGLOBAL_GET(std::vector<std::string>*,name)->push_back(arg);
    return true;
  }
  return false;
}

void CmdParser::LoadOptionFile()
{
  string opt_filename;
  if (!ZFLAG_opt_file.empty()) {
    opt_filename = ZFLAG_opt_file;
  } else {
    opt_filename = GetBase(zCurrentProgramName) + ".zopt";
  }

  if (FileExists(PathFile(InitialPath(), opt_filename))) {
    opt_filename = PathFile(InitialPath(), opt_filename);
  } else if (!ZFLAG_opt_dir.empty() && FileExists(PathFile(CompleteDirName(ZFLAG_opt_dir), opt_filename))) {
    opt_filename = PathFile(CompleteDirName(ZFLAG_opt_dir), opt_filename);
  } else {
    return;
  }

  ZLOGM << "Load option file from " << opt_filename << endl;

  ifstream fi(opt_filename);
  if (!fi.good()) {
    ZLOGE << "Option file exists but cannot open " << ZVAR(opt_filename) << endl;
    return;
  }
  YAML::Parser parser(fi);
  YAML::Node doc;
  parser.GetNextDocument(doc);

  // Decide what profile to use.
  // Load profile in use:
  // If in release, load profile in use_release:
  // If in debug, load profile in use_debug:
  // If no use/use_xxx, load "default"
  vector<string> profiles;
  for (YAML::Iterator it = doc.begin(); it != doc.end(); ++it) {
    std::string key, value;
    it.first() >> key;

    if (key == "use") {
      if (it.second().Type() == YAML::NodeType::Scalar) {
        it.second() >> value;
        profiles.push_back(value);
      } else {
        for (YAML::Iterator it2 = it.second().begin(); it2 != it.second().end(); ++it2) {
          *it2 >> value;
          profiles.push_back(value);
        }
      }
    }

#ifdef ZZZ_DEBUG
    if (key == "use_debug") {
#else
    if (key == "use_release") {
#endif
      it.second() >> value;
      profiles.push_back(value);
    }
  }

  for (vector<string>::iterator vi = profiles.begin(); vi != profiles.end(); ++vi) {
    if (const YAML::Node *p_node = doc.FindValue(*vi)) {
      ZLOGM << "Load option profile " << ZVAR2("profile", *vi) << endl;
      LoadOptionProfile(p_node);
    } else {
      ZLOGE << "Unable to find profile " << ZVAR2("profile", *vi) << endl;
    }
  }

  if (profiles.empty()) {
    if (const YAML::Node *p_node = doc.FindValue("default")) {
      LoadOptionProfile(p_node);
    }
  }

  fi.close();
}

void CmdParser::LoadOptionProfile(const YAML::Node *p_node) {
  for(YAML::Iterator it = p_node->begin(); it != p_node->end(); ++it) {
    std::string key, value;
    it.first() >> key;
    key = string("ZFLAG_") + key;
    if (it.second().Type() == YAML::NodeType::Scalar) {
      it.second() >> value;
      ZLOGD << ZVAR2("Flag", key) << ZVAR2("Value", value) <<endl;
      bool res = ParseLongOption(key, value.c_str());
      if (!res)
        ZLOGE << "Unknown option: " << ZVAR(key) << ZVAR(value) <<endl;
    } else {
      for (YAML::Iterator it2 = it.second().begin(); it2 != it.second().end(); ++it2) {
        *it2 >> value;
        ZLOGD << ZVAR2("Flag", key) << ZVAR2("Value", value) <<endl;
        bool res = ParseLongOption(key, value.c_str());
        if (!res)
          ZLOGE << "Unknown option: " << ZVAR(key) << ZVAR(value) <<endl;
      }
    }
  }
}

}  // namespace zzz

